import os
import json
import jieba
from pypinyin import Style, lazy_pinyin

def load_pinyin_cache():
    """加载拼音缓存文件，如果不存在则创建"""
    cache_path = os.path.join(os.path.dirname(__file__), "pinyin_cache.json")
    
    if not os.path.exists(cache_path):
        from collections import defaultdict
        print("正在构建全量拼音缓存，可能需要10-20分钟...")
        word_dict = defaultdict(list)
        reverse_dict = {}

        # 遍历CJK统一汉字区块（U+4E00-U+9FA5）
        for code_point in range(0x4E00, 0x9FA6):
            char = chr(code_point)
            try:
                # 获取带声调拼音
                py = ''.join(lazy_pinyin(char, style=Style.TONE))
                word_dict[py].append(char)
                reverse_dict[char] = py
            except:
                continue
        
        # 保存缓存
        with open(cache_path, 'w', encoding='utf-8') as f:
            json.dump({'forward': word_dict, 'reverse': reverse_dict}, f)

    with open(cache_path, 'r', encoding='utf-8') as f:
        return json.load(f)

def find_homophones(word):
    """生成任意同音同调组合"""
    from itertools import product
    try:
        # 分解输入为单字
        chars = list(word)
        
        # 获取每个字的所有同音字
        homophone_lists = []
        for c in chars:
            py = WORD_TO_PINYIN.get(c, ''.join(lazy_pinyin(c, style=Style.TONE)))
            homophone_lists.append(PINYIN_CACHE.get(py, [c]))
        
        # 生成笛卡尔积
        return [''.join(p) for p in product(*homophone_lists)]
    except:
        return []

# 修改缓存加载方式
cache_data = load_pinyin_cache()
PINYIN_CACHE = cache_data['forward']
WORD_TO_PINYIN = cache_data['reverse']

# 删除以下重复的函数定义（已注释掉）
# def find_homophones(word):
#    """查找任意词语的同音词（包含原词关系）"""
#    try:
#        key = WORD_TO_PINYIN.get(word, ''.join(lazy_pinyin(word, style=Style.TONE)))
#        return PINYIN_CACHE.get(key, [])
#    except:
#        return []

def process_phrase_group(phrases):
    """处理词组输入并构建反向索引"""
    reverse_map = {}  # 移除了多维数组相关代码
    
    for phrase in phrases:
        clean_phrase = phrase.strip()
        if not clean_phrase:
            continue
        
        homophones = find_homophones(clean_phrase)
        
        # 只构建反向映射关系
        for homophone in homophones:
            reverse_map.setdefault(homophone, []).append(clean_phrase)
            
    return reverse_map  # 仅返回反向索引

if __name__ == "__main__":
    while True:
        input_str = input("请输入逗号分隔的词组表（输入 q 退出）: ").strip()
        if input_str.lower() == 'q':
            break
            
        phrases = [p.strip() for p in input_str.split(',') if p.strip()]
        if not phrases:
            continue
            
        reverse_index = process_phrase_group(phrases)
        
        # 修改保存部分的代码
        with open('cmdword.py', 'w', encoding='utf-8') as f:
            f.write(f"# 自动生成的同音词索引\nindex = {repr(reverse_index)}\n\n")
            f.write("def check(word):\n    return index.get(word, [])\n")
        # --------------------------------------------------------------
        
        print(f"\n=== 反向索引原始数据 ===")
        print(reverse_index)
        print("======================")
        
        while True:
            query = input("输入要查找的同音词（回车返回主菜单）: ").strip()
            if not query:
                break
                
            originals = reverse_index.get(query, [])
            print(f"匹配的原词: {originals if originals else '无'}")